from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator
import os



'''
    precision: [T, R, K, A, M]  (iou, recall, cls, area range, max dets)
               [10, 101, 3, 4, 3]

    T 表示COCO计算时采用的10个IoU值, 从0.5到0.95每间隔0.05取一个值
    R 表示COCO计算时采用的每一个概率阈值，这里是从0到1每间隔0.01（即一个百分点）取一个值, 共101的值
    K 表示检测任务中检测的目标类别数，假设针对COCO数据集就为80
    A 表示检测任务中针对的目标尺度类型 共4个值:
        第一个表示没有限制, 
        第二个代表小目标（area < 32^2）
        第三个代表中等目标（32^2 < area < 96^2）
        第四个代表大目标（area > 96^2）
    M 表示每张图片最大检测目标个数，COCO中有1,10,100共3个值
'''

def _summarize(coco, ap=True, catId=None, iouThr=None, areaRng='all', maxDets=100):
    p = coco.params
    iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
    titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
    typeStr = '(AP)' if ap == 1 else '(AR)'
    iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
        if iouThr is None else '{:0.2f}'.format(iouThr)

    # areaRng ['all', 'small', 'medium', 'large']
    aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
    mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]

    if ap:
        # dimension of precision: [TxRxKxAxM]
        s = coco.eval['precision']
        # IoU
        if iouThr is not None:
            t = np.where(iouThr == p.iouThrs)[0]
            s = s[t]

        if isinstance(catId, int):
            s = s[:, :, catId, aind, mind]
        else:
            s = s[:, :, :, aind, mind]

    else:
        # dimension of recall: [TxKxAxM]
        s = coco.eval['recall']
        if iouThr is not None:
            t = np.where(iouThr == p.iouThrs)[0]
            s = s[t]

        if isinstance(catId, int):
            s = s[:, catId, aind, mind]
        else:
            s = s[:, :, aind, mind]

    if len(s[s > -1]) == 0:
        mean_s = -1
    else:
        mean_s = np.mean(s[s > -1])

    print_string = iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s)
    return mean_s, print_string

if __name__ == "__main__":
    label_json_path = './val_new.json'
    res_json_path = './test_fusion.json'
    categories_info = {
        1: '1 (rect_eye)',
        2: '2 (sphere_eye)',
        # 3: '3 (box_eye)',
    }


    coco_gt = COCO(label_json_path)
    coco_dt = coco_gt.loadRes(res_json_path)
    coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()

    mAP, _ = _summarize(coco_eval)
    print(mAP)